package burp;

import java.awt.Component;
import java.awt.Container;
import java.awt.FlowLayout;
import java.awt.GridLayout;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;

import javax.swing.JPanel;
import javax.swing.JScrollPane;
import javax.swing.JSplitPane;
import javax.swing.JTabbedPane;
import javax.swing.JTable;
import javax.swing.JButton;
import javax.swing.JTextField;
import javax.swing.JFrame;
import javax.swing.SwingUtilities;
import javax.swing.table.AbstractTableModel;
import javax.swing.table.TableModel;

import java.util.regex.*;
import java.io.PrintWriter;

import burp.UCManager;
public class BurpExtender extends AbstractTableModel implements IBurpExtender,
		ITab, IHttpListener, IMessageEditorController {
	private IBurpExtenderCallbacks callbacks;
	private IExtensionHelpers helpers;
	private JSplitPane splitPane;
	private IMessageEditor requestViewer;
	private IMessageEditor responseViewer;
	private final List<LogEntry> log = new ArrayList<LogEntry>();
	private IHttpRequestResponse currentlyDisplayedItem;
	private JTextField access_token;
	private JTextField mac_key;
	private JTextField uc_host;

	//
	// implement IBurpExtender
	//

	public void registerExtenderCallbacks(final IBurpExtenderCallbacks callbacks) {
		// keep a reference to our callbacks object
		this.callbacks = callbacks;

		// obtain an extension helpers object
		helpers = callbacks.getHelpers();

		// set our extension name
		callbacks.setExtensionName("UC Plugin");

		// create our UI
		SwingUtilities.invokeLater(new Runnable() {
			public void run() {

				// main split pane
				splitPane = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT);

				// table of log entries
				Table logTable = new Table(BurpExtender.this);
				JScrollPane scrollPane = new JScrollPane(logTable);
				// splitPane.setLeftComponent(scrollPane);

				JPanel contPanel = new JPanel(new GridLayout(2, 1));
				JPanel arg = new JPanel(new GridLayout(3, 1));
				access_token = new JTextField(128);
				access_token.setSize(4, 32);
				access_token.setText("复制ACCESS_TOKEN到此");
				mac_key = new JTextField(16);
				mac_key.setSize(4, 32);
				mac_key.setText("复制MAC_KEY到此");

				uc_host = new JTextField(32);
				uc_host.setSize(4, 32);
				uc_host.setText("101uccenter.dev.web.nd");

				contPanel.add(scrollPane);
				arg.add(access_token);
				arg.add(mac_key);
				arg.add(uc_host);
				contPanel.add(arg);
				splitPane.setLeftComponent(contPanel);

				// tabs with request/response viewers
				JTabbedPane tabs = new JTabbedPane();
				requestViewer = callbacks.createMessageEditor(
						BurpExtender.this, false);
				responseViewer = callbacks.createMessageEditor(
						BurpExtender.this, false);
				tabs.addTab("Request", requestViewer.getComponent());
				tabs.addTab("Response", responseViewer.getComponent());
				splitPane.setRightComponent(tabs);

				// customize our UI components
				callbacks.customizeUiComponent(splitPane);
				callbacks.customizeUiComponent(logTable);

				callbacks.customizeUiComponent(access_token);
				callbacks.customizeUiComponent(mac_key);
				callbacks.customizeUiComponent(uc_host);
				callbacks.customizeUiComponent(contPanel);
				callbacks.customizeUiComponent(scrollPane);
				callbacks.customizeUiComponent(tabs);

				// add the custom tab to Burp's UI
				callbacks.addSuiteTab(BurpExtender.this);

				// register ourselves as an HTTP listener
				callbacks.registerHttpListener(BurpExtender.this);
			}

		});
	}

	//
	// implement ITab
	//

	public String getTabCaption() {
		return "UC插件";
	}

	public Component getUiComponent() {
		return splitPane;
	}

	//
	// implement IHttpListener
	//

	public void processHttpMessage(int toolFlag, boolean messageIsRequest,
			IHttpRequestResponse messageInfo) {

		
		IHttpService httpService = messageInfo.getHttpService();
		byte[] req = messageInfo.getRequest();
		String r = new String(req);
		
		String pattern = "Host:.*\r\n";
		Pattern reg = Pattern.compile(pattern);
		Matcher m = reg.matcher(r);
		String req_host = "";
		if (m.find()) {
			String rst = m.group(0).trim(); 
			req_host = rst.split(":")[1].trim();
		}
		if(req_host == "") {
			req_host = httpService.getHost();
		}
		String host = uc_host.getText().trim();
		if (host.indexOf(req_host) != -1) {
			if (messageIsRequest) {
				messageInfo.setHighlight("red");
				
				int i = r.indexOf("/");
				int x = r.indexOf("HTTP/");
				String url = r.substring(i, x).trim();
				UCManager ucm = new UCManager();
				String id = access_token.getText().trim();
				String macKey = mac_key.getText().trim();
				String method = r.substring(0, i).trim();
				
				String auth = ucm.getAuth(id, req_host, method, url, macKey);
				auth = "Authorization: " + auth;
				
				
				if(r.indexOf("Authorization:") !=-1) {
					r = r.replaceAll("Authorization.*", auth);
					messageInfo.setComment("修改了Authorization");
				} else if(r.indexOf("authorization:") !=-1) {
					r = r.replaceAll("authorization.*", auth);
					messageInfo.setComment("修改了Authorization");
				} else {
					int y = r.indexOf("Host:");
					r = r.substring(0, y) + auth +"\n" + r.substring(y, r.length());
					messageInfo.setComment("添加了Authorization");
				}
				
				byte[] data = r.getBytes();
				messageInfo.setRequest(data);
				
			} else {
				// create a new log entry with the message details
				// only process responses
				synchronized (log) {
					int row = log.size();
					log.add(new LogEntry(toolFlag, callbacks
							.saveBuffersToTempFiles(messageInfo), helpers
							.analyzeRequest(messageInfo).getUrl()));
					fireTableRowsInserted(row, row);
				}
			}
		}

	}

	//
	// extend AbstractTableModel
	//

	public int getRowCount() {
		return log.size();
	}

	public int getColumnCount() {
		return 2;
	}

	@Override
	public String getColumnName(int columnIndex) {
		switch (columnIndex) {
		case 0:
			return "Tool";
		case 1:
			return "URL";
		default:
			return "";
		}
	}

	@Override
	public Class<?> getColumnClass(int columnIndex) {
		return String.class;
	}

	public Object getValueAt(int rowIndex, int columnIndex) {
		LogEntry logEntry = log.get(rowIndex);

		switch (columnIndex) {
		case 0:
			return callbacks.getToolName(logEntry.tool);
		case 1:
			return logEntry.url.toString();
		default:
			return "";
		}
	}

	//
	// implement IMessageEditorController
	// this allows our request/response viewers to obtain details about the
	// messages being displayed
	//

	public byte[] getRequest() {
		return currentlyDisplayedItem.getRequest();
	}

	public byte[] getResponse() {
		return currentlyDisplayedItem.getResponse();
	}

	public IHttpService getHttpService() {
		return currentlyDisplayedItem.getHttpService();
	}

	//
	// extend JTable to handle cell selection
	//

	private class Table extends JTable {
		public Table(TableModel tableModel) {
			super(tableModel);
		}

		@Override
		public void changeSelection(int row, int col, boolean toggle,
				boolean extend) {
			// show the log entry for the selected row
			LogEntry logEntry = log.get(row);
			requestViewer.setMessage(logEntry.requestResponse.getRequest(),
					true);
			responseViewer.setMessage(logEntry.requestResponse.getResponse(),
					false);
			currentlyDisplayedItem = logEntry.requestResponse;

			super.changeSelection(row, col, toggle, extend);
		}
	}

	//
	// class to hold details of each log entry
	//

	private static class LogEntry {
		final int tool;
		final IHttpRequestResponsePersisted requestResponse;
		final URL url;

		LogEntry(int tool, IHttpRequestResponsePersisted requestResponse,
				URL url) {
			this.tool = tool;
			this.requestResponse = requestResponse;
			this.url = url;
		}
	}
}
